package com.dahuyou.link.monitoring;

import java.util.*;

public class TrackManager {

    private static final ThreadLocal<Stack<String>> track = new ThreadLocal<Stack<String>>();
    private static final ThreadLocal<Map<String, List<Span>>> trackWithSpan = new ThreadLocal<>();

    private static Span createSpan(String className, String methodName) {
        Map<String, List<Span>> stack = trackWithSpan.get();
        String traceId = TrackContext.getTraceId();

        if (stack == null) {
            stack = new HashMap<>();
            trackWithSpan.set(stack);
        }
        if (!stack.containsKey(traceId)) {
            List<Span> spanList = new ArrayList<>();
            stack.put(traceId, spanList);
        }
        // 创建方法调用span
        Span span = new Span();
        span.setClassName(className);
        span.setMethodName(methodName);
        span.setTraceId(traceId);
        // 添加span到traceid对应的span调用集合中
        stack.get(traceId).add(span);
        return span;
    }

    private static String createSpan() {
        Stack<String> stack = track.get();
        if (stack == null) {
            stack = new Stack<>();
            track.set(stack);
        }
        String linkId;
        if (stack.isEmpty()) {
            linkId = TrackContext.getTraceId();
            if (linkId == null) {
                linkId = "nvl";
                TrackContext.setTraceId(linkId);
            }
        } else {
            linkId = stack.peek();
            TrackContext.setTraceId(linkId);
        }
        return linkId;
    }

    public static String createEntrySpan() {
        String span = createSpan();
        Stack<String> stack = track.get();
        stack.push(span);
        return span;
    }

    public static Span createEntrySpan(String className, String methodName) {
        Span span = createSpan(className, methodName);
        return span;
    }


    public static String getExitSpan() {
        Stack<String> stack = track.get();
        if (stack == null || stack.isEmpty()) {
            TrackContext.clear();
            return null;
        }
        return stack.pop();
    }

    public static String getCurrentSpan() {
        Stack<String> stack = track.get();
        if (stack == null || stack.isEmpty()) {
            return null;
        }
        return stack.peek();
    }

    public static void getAllSpan() {
        System.out.println(trackWithSpan.get().get(TrackContext.getTraceId()));
    }

    public static String getCurrentTraceId() {
        return TrackContext.getTraceId();
    }
}
